{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e3bfcca6-7d58-4df8-ae62-f5914b6d25fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import covalent as ct\n",
    "import numpy as np\n",
    "import time\n",
    "import random\n",
    "import platform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "14675748",
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_results = {}\n",
    "benchmark_results['platform'] = {}\n",
    "benchmark_results['platform']['arch'] = platform.architecture()\n",
    "benchmark_results['platform']['system'] = platform.system()\n",
    "benchmark_results['platform']['machine'] = platform.machine()\n",
    "benchmark_results['platform']['os'] = platform.node()\n",
    "benchmark_results['platform']['python_version'] = platform.python_version()\n",
    "\n",
    "# Setup for workflows\n",
    "benchmark_results['workflows'] = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "35877bf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.electron\n",
    "def add(x: int, y: int):\n",
    "    return x + y\n",
    "\n",
    "@ct.electron\n",
    "def multiply(x: int, y: int):\n",
    "    return x*y\n",
    "\n",
    "@ct.lattice\n",
    "def workflow(x: int, y: int):\n",
    "    r1 = add(x, y)\n",
    "    r2 = multiply(r1, y)\n",
    "    return r1\n",
    "\n",
    "res = ct.dispatch_sync(workflow)(2, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f97b7f7",
   "metadata": {},
   "source": [
    "#### Basic 2 node workflow (measure overhead of covalent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "68290320",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_label = 'simple_two_node_workflow'\n",
    "benchmark_results['workflows'][workflow_label] = {}\n",
    "\n",
    "@ct.electron\n",
    "def add(x: int, y: int):\n",
    "    return x + y\n",
    "\n",
    "@ct.electron\n",
    "def multiply(x: int, y: int):\n",
    "    return x*y\n",
    "\n",
    "@ct.lattice\n",
    "def workflow(x, y):\n",
    "    r1 = add(x, y)\n",
    "    r2 = multiply(r1, y)\n",
    "    return r1\n",
    "\n",
    "# Non covalent run\n",
    "benchmark_results['workflows'][workflow_label]['direct'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['direct']['duration'] = []\n",
    "for i in range(5):\n",
    "    start = time.time()\n",
    "    res = workflow(2, 3)\n",
    "    end = time.time()\n",
    "    benchmark_results['workflows'][workflow_label]['direct']['duration'].append(end-start)\n",
    "benchmark_results['workflows'][workflow_label]['direct']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['direct']['duration'])\n",
    "\n",
    "benchmark_results['workflows'][workflow_label]['covalent'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['duration'] = []\n",
    "\n",
    "for i in range(5):\n",
    "    result = ct.dispatch_sync(workflow)(2, 3)\n",
    "    duration = (result.end_time - result.start_time).total_seconds()\n",
    "    benchmark_results['workflows'][workflow_label]['covalent']['duration'].append(duration)\n",
    "\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['covalent']['duration'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04a03a28",
   "metadata": {},
   "source": [
    "#### Workflow with compute intensive electrons (primality testing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7d551a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_label = 'primality_tests'\n",
    "benchmark_results['workflows'][workflow_label] = {}\n",
    "\n",
    "@ct.electron\n",
    "def is_prime(n: int) -> bool:\n",
    "    \"\"\"Primality test using 6k+-1 optimization.\"\"\"\n",
    "    if n <= 3:\n",
    "        return n > 1\n",
    "    if not n%2 or not n%3:\n",
    "        return False\n",
    "    i = 5\n",
    "    stop = int(n**0.5)\n",
    "    while i <= stop:\n",
    "        if not n%i or not n%(i + 2):\n",
    "            return False\n",
    "        i += 6\n",
    "    return True\n",
    "\n",
    "# numbers to test\n",
    "nums_to_test = [random.randint(1000, 10000) for i in range(50)]\n",
    "\n",
    "@ct.lattice\n",
    "def primality_tests():\n",
    "    res = []\n",
    "    for i in nums_to_test:\n",
    "        entry = {}\n",
    "        entry['num'] = i\n",
    "        entry['is_prime'] = is_prime(i)\n",
    "        res.append(entry)\n",
    "    return res\n",
    "\n",
    "# Non covalent run\n",
    "benchmark_results['workflows'][workflow_label]['direct'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['direct']['duration'] = []\n",
    "for i in range(5):\n",
    "    start = time.time()\n",
    "    res = primality_tests()\n",
    "    end = time.time()\n",
    "    benchmark_results['workflows'][workflow_label]['direct']['duration'].append(end-start)\n",
    "benchmark_results['workflows'][workflow_label]['direct']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['direct']['duration'])\n",
    "\n",
    "# Covalent run\n",
    "benchmark_results['workflows'][workflow_label]['covalent'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['duration'] = []\n",
    "for i in range(5):\n",
    "    result = ct.dispatch_sync(primality_tests)()\n",
    "    duration = (result.end_time - result.start_time).total_seconds()\n",
    "    benchmark_results['workflows'][workflow_label]['covalent']['duration'].append(duration)\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['covalent']['duration'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9966515",
   "metadata": {},
   "source": [
    "### Memory intensive workflow (martix multiplication)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "87339f32",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_label = 'matrix_multiplication'\n",
    "benchmark_results['workflows'][workflow_label] = {}\n",
    "\n",
    "@ct.electron\n",
    "def create_matrix(arraysize: int):\n",
    "    return np.random.random((arraysize, arraysize))\n",
    "\n",
    "@ct.electron\n",
    "def matrix_multiply(a: np.ndarray, b: np.ndarray):\n",
    "    return np.matmul(a, b)\n",
    "\n",
    "@ct.lattice\n",
    "def matrix_multiplication():\n",
    "    for arraysize in [256, 512, 1024, 2048]:\n",
    "        a = create_matrix(arraysize)\n",
    "        b = create_matrix(arraysize)\n",
    "        matrix_multiply(a, b)\n",
    "\n",
    "# Non covalent run\n",
    "benchmark_results['workflows'][workflow_label]['direct'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['direct']['duration'] = []\n",
    "for i in range(5):\n",
    "    start = time.time()\n",
    "    res = matrix_multiplication()\n",
    "    end = time.time()\n",
    "    benchmark_results['workflows'][workflow_label]['direct']['duration'].append(end-start)\n",
    "benchmark_results['workflows'][workflow_label]['direct']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['direct']['duration'])\n",
    "\n",
    "# Covalent run\n",
    "benchmark_results['workflows'][workflow_label]['covalent'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['duration'] = []\n",
    "for i in range(5):\n",
    "    result = ct.dispatch_sync(matrix_multiplication)()\n",
    "    duration = (result.end_time - result.start_time).total_seconds()\n",
    "    benchmark_results['workflows'][workflow_label]['covalent']['duration'].append(duration)\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['covalent']['duration'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b159411",
   "metadata": {},
   "source": [
    "### Workflow with chain dependencies (simple electrons)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a20d2fe6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "workflow_label = 'linear_add_workflow'\n",
    "benchmark_results['workflows'][workflow_label] = {}\n",
    "\n",
    "@ct.electron\n",
    "def add(x: int, y: int):\n",
    "    return x + y\n",
    "\n",
    "@ct.lattice\n",
    "def add_workflow(N: int):\n",
    "    for i in range(N):\n",
    "        if i == 0:\n",
    "            r1 = add(1, 1)\n",
    "        else:\n",
    "            r1 = add(r1, 1)\n",
    "    return r1\n",
    "\n",
    "# Non covalent run\n",
    "benchmark_results['workflows'][workflow_label]['direct'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['direct']['duration'] = []\n",
    "for i in range(5):\n",
    "    start = time.time()\n",
    "    res = add_workflow(50)\n",
    "    end = time.time()\n",
    "    benchmark_results['workflows'][workflow_label]['direct']['duration'].append(end-start)\n",
    "benchmark_results['workflows'][workflow_label]['direct']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['direct']['duration'])\n",
    "\n",
    "# Covalent run\n",
    "benchmark_results['workflows'][workflow_label]['covalent'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['duration'] = []\n",
    "for i in range(5):\n",
    "    result = ct.dispatch_sync(add_workflow)(50)\n",
    "    duration = (result.end_time - result.start_time).total_seconds()\n",
    "    benchmark_results['workflows'][workflow_label]['covalent']['duration'].append(duration)\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['covalent']['duration'])\n",
    "\n",
    "workflow_label = 'linear_multiply_workflow'\n",
    "benchmark_results['workflows'][workflow_label] = {}\n",
    "\n",
    "@ct.electron\n",
    "def multiply(x: int, y: int):\n",
    "    return x*y\n",
    "\n",
    "# Result can potentially overflow\n",
    "@ct.lattice\n",
    "def multiply_workflow(N: int):\n",
    "    for i in range(N):\n",
    "        if i == 0:\n",
    "            r1 = multiply(2, 1)\n",
    "        else:\n",
    "            r1 = multiply(r1, 2)\n",
    "    return r1\n",
    "\n",
    "# Non covalent run\n",
    "benchmark_results['workflows'][workflow_label]['direct'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['direct']['duration'] = []\n",
    "for i in range(5):\n",
    "    start = time.time()\n",
    "    res = multiply_workflow(50)\n",
    "    end = time.time()\n",
    "    benchmark_results['workflows'][workflow_label]['direct']['duration'].append(end-start)\n",
    "benchmark_results['workflows'][workflow_label]['direct']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['direct']['duration'])\n",
    "\n",
    "# Covalent run\n",
    "benchmark_results['workflows'][workflow_label]['covalent'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['duration'] = []\n",
    "for i in range(5):\n",
    "    result = ct.dispatch_sync(multiply_workflow)(50)\n",
    "    duration = (result.end_time - result.start_time).total_seconds()\n",
    "    benchmark_results['workflows'][workflow_label]['covalent']['duration'].append(duration)\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['covalent']['duration'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3be43ca1",
   "metadata": {},
   "source": [
    "### Embarassingly parallel workload (fileio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0807b261",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_label = 'parallel_fileio'\n",
    "benchmark_results['workflows'][workflow_label] = {}\n",
    "\n",
    "import tempfile\n",
    "import secrets\n",
    "\n",
    "@ct.electron\n",
    "def create_delete_tempfile():\n",
    "    fp = tempfile.NamedTemporaryFile(delete=True)\n",
    "    # thousand lines per file\n",
    "    for i in range(1000):\n",
    "        fp.write(secrets.token_bytes(16384))\n",
    "    fp.close()\n",
    "\n",
    "@ct.lattice\n",
    "def parallel_fileio(N: int):\n",
    "    for i in range(N):\n",
    "        create_delete_tempfile()\n",
    "\n",
    "# Non covalent run\n",
    "benchmark_results['workflows'][workflow_label]['direct'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['direct']['duration'] = []\n",
    "for i in range(5):\n",
    "    start = time.time()\n",
    "    res = parallel_fileio(50)\n",
    "    end = time.time()\n",
    "    benchmark_results['workflows'][workflow_label]['direct']['duration'].append(end-start)\n",
    "benchmark_results['workflows'][workflow_label]['direct']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['direct']['duration'])\n",
    "\n",
    "# Covalent run\n",
    "benchmark_results['workflows'][workflow_label]['covalent'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['duration'] = []\n",
    "for i in range(5):\n",
    "    result = ct.dispatch_sync(parallel_fileio)(50)\n",
    "    duration = (result.end_time - result.start_time).total_seconds()\n",
    "    benchmark_results['workflows'][workflow_label]['covalent']['duration'].append(duration)\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['covalent']['duration'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8b343af",
   "metadata": {},
   "source": [
    "### Sublattices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e01c5c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_label = 'sublattices'\n",
    "benchmark_results['workflows'][workflow_label] = {}\n",
    "\n",
    "@ct.electron\n",
    "@ct.lattice\n",
    "def sublattice(x: int, y: int):\n",
    "    r1 = add(x, y)\n",
    "    r2 = multiply(r1, y)\n",
    "    return r1 + r2\n",
    "\n",
    "@ct.lattice\n",
    "def sublattice_workflow(num_tasks: int):\n",
    "    results = []\n",
    "    for i in range(num_tasks):\n",
    "        results.append(sublattice(i, i+2))\n",
    "    return results\n",
    "\n",
    "# Non covalent run\n",
    "benchmark_results['workflows'][workflow_label]['direct'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['direct']['duration'] = []\n",
    "for i in range(5):\n",
    "    start = time.time()\n",
    "    res = sublattice_workflow(10)\n",
    "    end = time.time()\n",
    "    benchmark_results['workflows'][workflow_label]['direct']['duration'].append(end-start)\n",
    "benchmark_results['workflows'][workflow_label]['direct']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['direct']['duration'])\n",
    "\n",
    "# Covalent run\n",
    "benchmark_results['workflows'][workflow_label]['covalent'] = {}\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['duration'] = []\n",
    "for i in range(5):\n",
    "    result = ct.dispatch_sync(sublattice_workflow)(10)\n",
    "    duration = (result.end_time - result.start_time).total_seconds()\n",
    "    benchmark_results['workflows'][workflow_label]['covalent']['duration'].append(duration)\n",
    "benchmark_results['workflows'][workflow_label]['covalent']['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['covalent']['duration'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d47f6a60",
   "metadata": {},
   "source": [
    "#### Scaling analysis for sample compute intensive workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fc3fc97b",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_label = 'compute_intensive_scaling'\n",
    "benchmark_results['workflows'][workflow_label] = {}\n",
    "\n",
    "@ct.electron\n",
    "def is_prime(n: int) -> bool:\n",
    "    \"\"\"Primality test using 6k+-1 optimization.\"\"\"\n",
    "    if n <= 3:\n",
    "        return n > 1\n",
    "    if not n%2 or not n%3:\n",
    "        return False\n",
    "    i = 5\n",
    "    stop = int(n**0.5)\n",
    "    while i <= stop:\n",
    "        if not n%i or not n%(i + 2):\n",
    "            return False\n",
    "        i += 6\n",
    "    return True\n",
    "    \n",
    "# Measure scaling of the covalent execution\n",
    "@ct.lattice\n",
    "def primality_tests_scaling(num_nodes: int):\n",
    "    nums_to_test = [random.randint(1000, 10000) for i in range(num_nodes)]\n",
    "    res = []\n",
    "    for i in nums_to_test:\n",
    "        entry = {}\n",
    "        entry['num'] = i\n",
    "        entry['is_prime'] = is_prime(i)\n",
    "        res.append(entry)\n",
    "    return res\n",
    "\n",
    "# Covalent scaling\n",
    "benchmark_results['workflows'][workflow_label]['scaling'] = {}\n",
    "for j in [2, 4, 8, 16, 32, 64]:\n",
    "    benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"] = {}\n",
    "    benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"]['duration'] = []\n",
    "    for i in range(5):\n",
    "        result = ct.dispatch_sync(primality_tests_scaling)(j)\n",
    "        duration = (result.end_time - result.start_time).total_seconds()\n",
    "        benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"]['duration'].append(duration)\n",
    "    benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"]['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"]['duration'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e49b93c",
   "metadata": {},
   "source": [
    "#### Scaling analysis for parallel file i/o tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d052c20b",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_label = 'parallel_fileio_scaling'\n",
    "benchmark_results['workflows'][workflow_label] = {}\n",
    "\n",
    "import tempfile\n",
    "import secrets\n",
    "\n",
    "@ct.electron\n",
    "def create_delete_tempfile():\n",
    "    fp = tempfile.NamedTemporaryFile(delete=True)\n",
    "    # thousand lines per file\n",
    "    for i in range(1000):\n",
    "        fp.write(secrets.token_bytes(16384))\n",
    "    fp.close()\n",
    "\n",
    "@ct.lattice\n",
    "def parallel_fileio(N: int):\n",
    "    for i in range(N):\n",
    "        create_delete_tempfile()\n",
    "\n",
    "# Covalent run scaling\n",
    "benchmark_results['workflows'][workflow_label]['scaling'] = {}\n",
    "for j in [2, 4, 8, 16, 32, 64]:\n",
    "    benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"] = {}\n",
    "    benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"]['duration'] = []\n",
    "    for i in range(5):\n",
    "        result = ct.dispatch_sync(parallel_fileio)(j)\n",
    "        duration = (result.end_time - result.start_time).total_seconds()\n",
    "        benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"]['duration'].append(duration)\n",
    "    benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"]['avg_duration'] = np.mean(benchmark_results['workflows'][workflow_label]['scaling'][f\"{j}\"]['duration'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22522cde",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "96665897554d405553c080f7a3b66ad582de30e62526c6c9d18f08acffe92e34"
  },
  "kernelspec": {
   "display_name": "Python 3.8.13 ('covalent-qa')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
